// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <array>
#include <cmath>
#include <vector>
#include <melon/container/f14_map.h>
#include <melon/container/f14_map.h>
#include <melon/container/f14_set.h>

namespace kumo::pollux {
    /// Custom comparator and hash functors for floating point types. These are
    /// designed to ensure consistent NaN handling according to the following rules:
    /// -> NaN == NaN returns true, even for NaNs with differing binary
    ///   representations.
    /// -> NaN is considered greater than infinity. There is no ordering between
    ///    NaNs with differing binary representations.
    /// These can be passed to standard containers and functions like std::map and
    /// std::sort, etc.
    namespace util::floating_point {
        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareEquals {
            bool operator()(const FLOAT &lhs, const FLOAT &rhs) const {
                if (std::isnan(lhs) && std::isnan(rhs)) {
                    return true;
                }
                return lhs == rhs;
            }
        };

        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareLessThan {
            bool operator()(const FLOAT &lhs, const FLOAT &rhs) const {
                if (!std::isnan(lhs) && std::isnan(rhs)) {
                    return true;
                }
                return lhs < rhs;
            }
        };

        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareLessThanEqual {
            bool operator()(const FLOAT &lhs, const FLOAT &rhs) const {
                if (std::isnan(rhs)) {
                    return true;
                }
                return lhs <= rhs;
            }
        };

        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareGreaterThan {
            bool operator()(const FLOAT &lhs, const FLOAT &rhs) const {
                if (std::isnan(lhs) && !std::isnan(rhs)) {
                    return true;
                }
                return lhs > rhs;
            }
        };

        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareGreaterThanEqual {
            bool operator()(const FLOAT &lhs, const FLOAT &rhs) const {
                if (std::isnan(lhs)) {
                    return true;
                }
                return lhs >= rhs;
            }
        };

        template<
            typename FLOAT,
            std::enable_if_t<std::is_floating_point<FLOAT>::value, bool>  = true>
        struct NaNAwareHash {
            std::size_t operator()(const FLOAT &val) const noexcept {
                static const std::size_t kNanHash =
                        melon::hasher<FLOAT>{}(std::numeric_limits<FLOAT>::quiet_NaN());
                if (std::isnan(val)) {
                    return kNanHash;
                }
                return melon::hasher<FLOAT>{}(val);
            }
        };

        // Utility struct to provide a clean way of defining a hash set and map type
        // using melon::F14FastSet and melon::F14FastMap respectively with overrides for
        // floating point types.
        template<typename Key>
        class HashSetNaNAware : public melon::F14FastSet<Key> {
        };

        template<>
        class HashSetNaNAware<float>
                : public melon::F14FastSet<float, NaNAwareHash<float>, NaNAwareEquals<float> > {
        };

        template<>
        class HashSetNaNAware<double>
                : public melon::F14FastSet<double, NaNAwareHash<double>, NaNAwareEquals<double> > {
        };

        template<
            typename Key,
            typename Mapped,
            typename Alloc = melon::f14::DefaultAlloc<std::pair<Key const, Mapped> > >
        struct HashMapNaNAwareTypeTraits {
            using Type = melon::F14FastMap<
                Key,
                Mapped,
                melon::f14::DefaultHasher<Key>,
                melon::f14::DefaultKeyEqual<Key>,
                Alloc>;
        };

        template<typename Mapped, typename Alloc>
        struct HashMapNaNAwareTypeTraits<float, Mapped, Alloc> {
            using Type = melon::F14FastMap<
                float,
                Mapped,
                NaNAwareHash<float>,
                NaNAwareEquals<float>,
                Alloc>;
        };

        template<typename Mapped, typename Alloc>
        struct HashMapNaNAwareTypeTraits<double, Mapped, Alloc> {
            using Type = melon::F14FastMap<
                double,
                Mapped,
                NaNAwareHash<double>,
                NaNAwareEquals<double>,
                Alloc>;
        };
    } // namespace util::floating_point

    /// A static class that holds helper functions for DOUBLE type.
    class DoubleUtil {
    public:
        static const std::array<double, 309> kPowersOfTen;
    }; // DoubleUtil
} // namespace kumo::pollux
